from flask import Blueprint, request, jsonify, g, Response
from app.models import ChatSession, ChatMessage, SystemConfig, AIServer, AIModel, DoesNotExist
from app.utils import login_required, api_key_required, async_ai_stream,auto_api_key
import json

import logging
# 在 chat.py 顶部追加
import os
from werkzeug.utils import secure_filename
from app.test_ollama_v0 import stream_ollama

ALLOWED_EXT = {"txt", "py", "md", "pdf", "json"}   # 需要可自行扩展
chat_bp = Blueprint('chat', __name__)


@chat_bp.route('/sessions', methods=['GET'])
@login_required
async def get_chat_sessions():
    """获取用户的聊天会话列表"""
    sessions = (ChatSession
               .select()
               .where(ChatSession.user == g.user)
               .order_by(ChatSession.updated_at.desc()))
    
    return jsonify([{
        'id': s.id,
        'title': s.title,
        'created_at': s.created_at.isoformat(),
        'updated_at': s.updated_at.isoformat()
    } for s in sessions])

@chat_bp.route('/sessions/<int:currsession_id>', methods=['POST','DELETE'])
@login_required
async def delete_chat_session(currsession_id):
    """删除聊天记录"""
    try:
        chat_session:ChatSession = ChatSession.get_by_id(currsession_id)
        chat_session.delete_instance()

        return jsonify({
            'success':True
        }), 200
    except Exception as e:
        return jsonify({'error': f'创建会话失败：{str(e)}'}), 500

@chat_bp.route('/sessions/<int:currsession_id>/messages', methods=['POST','DELETE'])
@login_required
async def clear_chat_session(currsession_id):
    """清空聊天窗口"""
    chat_session = ChatSession.get_by_id(currsession_id)
    
    try:
        ChatMessage.delete().where(ChatMessage.session == chat_session).execute()

        return jsonify({
            'success':True
        }), 200
    except Exception as e:
        return jsonify({'error': f'创建会话失败：{str(e)}'}), 500


@chat_bp.route('/sessions', methods=['POST'])
@login_required
async def create_chat_session():
    """创建新的聊天会话"""
    data = request.get_json() or {}
    title = data.get('title', '新话题')
    
    try:
        session = ChatSession.create(
            user=g.user,
            title=title
        )
        return jsonify({
            'id': session.id,
            'title': session.title,
            'created_at': session.created_at.isoformat()
        }), 201
    except Exception as e:
        return jsonify({'error': f'创建会话失败：{str(e)}'}), 500

@chat_bp.route('/sessions/messages', methods=['POST'])
@login_required
async def get_session_messages():
    """获取指定会话的消息"""
    session_id  = request.form.get('session_id','')
    print(session_id)
    try:
        # 验证会话所有权
        session = ChatSession.get(
            (ChatSession.id == session_id) & 
            (ChatSession.user == g.user)
        )
        
        messages = (ChatMessage
                   .select()
                   .where(ChatMessage.session == session)
                   .order_by(ChatMessage.created_at.asc()))
        
        return jsonify([{
            'id': m.id,
            'role': m.role,
            'content': m.content,
            'created_at': m.created_at.isoformat()
        } for m in messages])
    except DoesNotExist:
        return jsonify({'error': '会话不存在或无权访问'}), 404

@chat_bp.route('/sessions/title', methods=['POST','PUT'])
@login_required
async def update_session_title():
    """更新会话标题"""
    data = request.get_json() or {}
    session_id = data.get('session_id',0)
    new_title = data.get('new_title','')

    if not new_title:
        return jsonify({'error': '标题不能为空'}), 400
    
    try:
        session = ChatSession.get(
            (ChatSession.id == session_id) & 
            (ChatSession.user == g.user)
        )
        session.title = new_title
        session.save()
        return jsonify({'status': 'success', 'title': session.title})
    except DoesNotExist:
        return jsonify({'error': '会话不存在或无权访问'}), 404

@chat_bp.route('/messages/save', methods=['POST'])
@login_required
async def save_ai_message():
    """保存AI消息"""
    data = request.get_json() or {}
    if not all(k in data for k in ['session_id', 'content']):
        return jsonify({'error': '会话ID和消息内容为必填项'}), 400
    
    try:
        session = ChatSession.get(
            (ChatSession.id == data['session_id']) & 
            (ChatSession.user == g.user)
        )
        
        message = ChatMessage.create(
            session=session,
            role='assistant',
            content=data['content']
        )
        
        return jsonify({
            'status': 'success',
            'message_id': message.id
        })
    except DoesNotExist:
        return jsonify({'error': '会话不存在或无权访问'}), 404



def allowed_file(filename):
    return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXT

def extract_text(file_stream, filename: str) -> str:
    """简单提取文本：pdf 用 PyPDF2，其余按 utf-8 读"""
    if filename.lower().endswith(".pdf"):
        try:
            from PyPDF2 import PdfReader
            reader = PdfReader(file_stream)
            return "\n".join(page.extract_text() or "" for page in reader.pages)
        except Exception:
            return "[PDF 解析失败]"
    else:
        return file_stream.read().decode("utf-8", errors="ignore")

# ========== 新路由：上传附件 ==========
@chat_bp.route("/upload", methods=["POST"])
@login_required
def upload_attachment():
    if "file" not in request.files:
        return jsonify({"error": "无文件域"}), 400
    file = request.files["file"]
    if file.filename == "":
        return jsonify({"error": "未选中文件"}), 400
    if not allowed_file(file.filename):
        return jsonify({"error": "不支持的文件类型"}), 400

    try:
        text = extract_text(file, secure_filename(file.filename))
        return jsonify({"text": text})
    except Exception as e:
        return jsonify({"error": f"解析失败: {str(e)}"}), 500

# ========== 改造原流式路由：同步 + 附件 ==========
@chat_bp.route("/messages/stream", methods=["GET","POST"])
@login_required
@auto_api_key
def stream_chat():

    data = request.get_json() or {}
    session_id = data.get('session_id')
    user_input = data.get('message', '').strip()
    attachment = data.get('attachment', '').strip()

    if not session_id or not user_input:
        return Response("data: 缺少 session_id 或 message\n\n", mimetype='text/event-stream')

    # 1. 查会话
    try:
        session = ChatSession.get((ChatSession.id == session_id) &
                                  (ChatSession.user == g.user))
    except DoesNotExist:
        return jsonify({"error": "会话不存在"}), 404

    # 2. 拼消息
    history = [{"role": m.role, "content": m.content}
               for m in ChatMessage.select()
               .where(ChatMessage.session == session)
               .order_by(ChatMessage.created_at.asc())]
    user_content = user_input
    if attachment:
        user_content = f"用户问题：{user_input}\n\n附件内容：\n{attachment}"
    history.append({"role": "user", "content": user_content})

    # 3. 保存用户消息（含附件提示）
    ChatMessage.create(session=session, role="user", content=user_content)

    # 4. 同步流式返回
    def generate():
        try:
            assistant_text = ""
            for delta in stream_ollama(history):
                assistant_text += delta
                yield json.dumps(dict(content=delta))
            # 正常结束
            ChatMessage.create(session=session, role="assistant", content=assistant_text)
            yield json.dumps(dict(content="[DONE]"))

        except Exception as e:
            yield json.dumps({"content": str(e)})
            yield json.dumps(dict(content="[DONE]"))

    return Response(generate(), mimetype="text/event-stream",
                    headers={"Cache-Control": "no-cache",
                             "X-Accel-Buffering": "no"})